#include "runner.h"
#include "types.h"
#include "constants.h"
#include "logger.h"
#include "configure.h"

#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <cstdlib>

#ifdef __i386
#define REG_SYSCALL orig_eax
#define REG_RET eax
#define REG_ARG0 ebx
#define REG_ARG1 ecx
#else
#define REG_SYSCALL orig_rax
#define REG_RET rax
#define REG_ARG0 rdi
#define REG_ARG1 rsi
#endif

Runner::Runner(int sol_id, int rid, const char *dir, const char *fname, char assign_symbol)
    : solution_id(sol_id), runner_id(rid)
{
    Configure conf(dir, fname, assign_symbol);
    init_data(conf);
}

Runner::Runner(int sol_id, int rid, const char *fname, char assign_symbol)
    : solution_id(sol_id), runner_id(rid)
{
    Configure conf(fname, assign_symbol);
    init_data(conf);
}

void
Runner::init_data(const Configure &conf)
{
    char buffer[BUFFER_SIZE];
    sprintf(buffer, "%s/run%d", conf["judger_home"], runner_id);

    result = JRT_AC;

    home_dir = buffer;
    hostname = conf["hostname"];
    username = conf["username"];
    passwd = conf["password"];
    dbname = conf["dbname"];
    port = conf["port"];

    init_mysql_conn();
    get_solution_info();
    get_problem_info();
    get_solution();
}

bool
Runner::init_mysql_conn()
{
    conn = mysql_init(NULL);
    const char timeout = 30;
    mysql_options(conn, MYSQL_OPT_CONNECT_TIMEOUT, &timeout);

    if (!mysql_real_connect(conn, hostname.c_str(), username.c_str(), passwd.c_str(), dbname.c_str(), atoi(port.c_str()), 0, 0)) {
        LOG(log, write_fatal, mysql_error(conn));
        return false;
    }
    const char *utf8sql = "set names utf8";
    if (mysql_real_query(conn, utf8sql, strlen(utf8sql))) {
        LOG(log, write_fatal, mysql_error(conn));
        return false;
    }
    return true;
}

void
Runner::get_solution_info()
{
    MYSQL_RES *res;
    MYSQL_ROW row; 

    char sql[BUFFER_SIZE];
    sprintf(sql, "SELECT `problem_id`, `user_id`, `language` FROM solution where `solution_id`=%d", solution_id);

    mysql_real_query(conn, sql, strlen(sql));
    res = mysql_store_result(conn);
    row = mysql_fetch_row(res);

    problem_id = atoi(row[0]);
    userid = row[1];
    lang = LANG_TYPE(atoi(row[2]));
    mysql_free_result(res);
}

void
Runner::get_problem_info()
{
    MYSQL_RES *res;
    MYSQL_ROW row; 

    char sql[BUFFER_SIZE];
    sprintf(sql, "SELECT `time_limit`,`memory_limit`,`spj` FROM problem where `problem_id`=%d", problem_id);

    mysql_real_query(conn, sql, strlen(sql));
    res = mysql_store_result(conn);
    row = mysql_fetch_row(res);

    time_lmt = atoi(row[0]);
    mem_lmt = atoi(row[1]);
    spj = atoi(row[2]) != 0;
    mysql_free_result(res);
}

void
Runner::get_solution()
{
    MYSQL_RES *res;
    MYSQL_ROW row; 

    char sql[BUFFER_SIZE], src_pth[BUFFER_SIZE];
    sprintf(sql, "SELECT `source` FROM source_code WHERE `solution_id`=%d", solution_id);

    mysql_real_query(conn, sql, strlen(sql));
    res = mysql_store_result(conn);
    row = mysql_fetch_row(res);

    sprintf(src_pth, "Solution.%s", LANG_EXT[lang]);
    FILE *fp_src = fopen(src_pth, "w");
    fprintf(fp_src, "%s", row[0]);
    fclose(fp_src);

    mysql_free_result(res);
}

void
Runner::update_solution(int result, int time, int memory)
{
    char sql[BUFFER_SIZE];
    sprintf(sql, "UPDATE solution SET `result`=%d,`time`=%d,`memory`=%d,`judgetime`=NOW() WHERE `solution_id`=%d LIMIT 1", result, time, memory, solution_id);

    mysql_real_query(conn, sql, strlen(sql));
}

void
Runner::insert_ceinfo()
{
    char sql[MAX_COMPILE_INFO_STORE << 1];
    char ceinfo[MAX_COMPILE_INFO_STORE], t[MAX_COMPILE_INFO_STORE];

    sprintf(sql, "DELETE FROM compileinfo WHERE `solution_id`=%d", solution_id);
    mysql_real_query(conn, sql, strlen(sql));

    int fd = open(CE_OUTPUT_FILE, O_RDONLY);
    read(fd, t, MAX_COMPILE_INFO_STORE - 1);
    close(fd);

    /* unsigned long mysql_real_escape_string(MYSQL *mysql, char *to, const char *from, unsigned long length) */
    mysql_real_escape_string(conn, ceinfo, t, strlen(t));

    sprintf(sql, "INSERT INTO compileinfo VALUES('%d','%s')", solution_id, ceinfo);
    mysql_real_query(conn, sql, strlen(sql));
}

void
Runner::update_user()
{
    char sql[BUFFER_SIZE];

    sprintf(sql, "UPDATE `users` SET `solved`=(SELECT count(DISTINCT `problem_id`) FROM `solution` WHERE `user_id`=\'%s\' AND `result`=\'%d\') WHERE `user_id`=\'%s\'", userid.c_str(), JRT_AC, userid.c_str()); 
    mysql_real_query(conn, sql, strlen(sql));

    sprintf(sql, "UPDATE `users` SET `submit`=(SELECT count(*) FROM `solution` WHERE `user_id`=\'%s\') WHERE `user_id`=\'%s\'", userid.c_str(), userid.c_str());
    mysql_real_query(conn, sql, strlen(sql));
}

void
update_problem()
{
    char sql[BUFFER_SIZE];

    sprintf(sql, "UPDATE `problem` SET `accepted`=(SELECT count(*) FROM `solution` WHERE `problem_id`=\'%d\' AND `result`=\'%d\') WHERE `problem_id`=\'%d\'", problem_id, JRT_AC, problem_id);
    mysql_real_query(conn, sql, strlen(sql));

    sprintf(sql, "UPDATE `problem` SET `submit`=(SELECT count(*) FROM `solution` WHERE `problem_id`=\'%d\') WHERE `problem_id`=\'%d\'", problem_id, problem_id);
    mysql_real_query(conn, sql, strlen(sql));
}

int
Runner::compile()
{
    int pid = fork();
    if (pid == 0) {
        struct rlimit rlim;
        rlim.rlim_max = COMPILE_MAX_CPU;
        rlim.rlim_cur = COMPILE_SOFT_CPU;
        setrlimit(RLIMIT_CPU, &rlim);

        rlim.rlim_max = COMPILE_MAX_FSIZE;
        rlim.rlim_cur = COMPILE_SOFT_FSIZE;
        setrlimit(RLIMIT_FSIZE, &rlim);

        rlim.rlim_max =  COMPILE_MAX_AS;
        rlim.rlim_cur =  COMPILE_SOFT_AS;
        setrlimit(RLIMIT_AS, &rlim);

        freopen(CE_OUTPUT_FILE, "w", stderr);

        switch (lang) {
        case LT_C:    execvp(CP_C[0], (char * const *) CP_C); break;
        case LT_CPP:  execvp(CP_CPP[0], (char * const *) CP_CPP); break;
        default:      LOG(log, write_warning, "Compile file's language type not supported!"); break;
        }
        exit(0);

    } else {
        int status = 0;
        waitpid(pid, &status, 0);
        return get_file_size(CE_OUTPUT_FILE);
    }
}

void
Runner::run_solution()
{
    nice(19);
    chdir(home_dir.c_str());
    
    freopen(INPUT_DATA_FILE, "r", stdin);
    freopen(USER_OUTPUT_FILE, "w", stdout);
    freopen(USER_ERROR_FILE, "a+", stderr);

    ptrace(PTRACE_TRACEME, 0, NULL, NULL);

    if (lang != LT_JAVA) chroot(home_dir.c_str());

    while (setgid(1536) != 0) sleep(1);
    while (setuid(1536) != 0) sleep(1);
    while (setresuid(1536, 1536, 1536) != 0) sleep(1);

    struct rlimit rlim;
    rlim.rlim_cur = (time_lmt - used_time / 1000) + 1;
    rlim.rlim_max = rlim.rlim_cur;
    setrlimit(RLIMIT_CPU, &rlim);

    alarm(0);
    alarm(time_lmt * 10);
  
    rlim.rlim_max = STD_F_LIM + STD_MB;
    rlim.rlim_cur = STD_F_LIM;
    setrlimit(RLIMIT_FSIZE, &rlim);

    rlim.rlim_cur = 1;
    rlim.rlim_max = 1;
    setrlimit(RLIMIT_NPROC, &rlim);

    rlim.rlim_cur = STD_MB << 6;
    rlim.rlim_max = STD_MB << 6;
    setrlimit(RLIMIT_STACK, &rlim);

    rlim.rlim_cur = STD_MB * mem_lmt / 2 * 3;
    rlim.rlim_max = STD_MB * mem_lmt * 2;
    setrlimit(RLIMIT_AS, &rlim);

    switch (lang) {
    case LT_C:
    case LT_CPP:
        execl("./Solution", "./Solution", (char *)NULL);
        break;
    case LT_PYTHON:
        execl("/python", "/python", "Solution.py", (char *)NULL);
        break;
    default:
        LOG(log, write_warning, "File's language type not supported!");
        break;
    }
    exit(0);
}

int
Runner::get_proc_status(int pid, const char *mark)
{
    int ret;
    char path[BUFFER_SIZE];
    sprintf(path, "/proc/%d/status", pid);

    Configure proc_conf(path, ':');
    sscanf(proc_conf[mark], "%d", &ret);

    return ret;
}

void
Runner::write_runtime_err(char *err)
{
    FILE *ferr = fopen(USER_ERROR_FILE, "a+");

    fprintf(ferr, "Runtime Error : [%s]\n", err);
    fclose(ferr);
}

void
watch_solution(pid_t pid)
{
    int status, sig, exitcode;
    struct user_regs_struct reg;
    struct rusage ruse;
    int max_mem, sub = 0;

    while (true) {
        wait4(pid, &status, 0, &ruse);
        max_mem = get_proc_status(pid, "VmPeak") << 10;

        if (max_men > used_mem) used_mem = max_mem;
        if (used_mem > mem_lmt * STD_MB) {
            if (result == JRT_AC) result = JRT_MLE;
            ptrace(PTRACE_KILL, pid, NULL, NULL);
            break;
        }
        if (WIFEXITED(status)) break;
        if (get_file_size(USER_ERROR_FILE) > 0) {
            result = JRT_RE;
            ptrace(PTRACE_KILL, pid, NULL, NULL);
            break;
        }
        if (!spj && get_file_size(USER_OUTPUT_FILE) > get_file_size(OUTPUT_DATA_FILE) * 2 + 1024) {
            result = JRT_OLE;
            ptrace(PTRACE_KILL, pid, NULL, NULL);
            break;
        }

        exitcode = WEXITSTATUS(status);
        if (exitcode != 0) {
            if (result == JRT_AC) {
                switch (exitcode) {
                    case SIGCHLD:
                    case SIGALRM: alarm(0);
                    case SIGKILL:
                    case SIGXCPU: result = JRT_TLE; break;
                    case SIGXFSZ: result = JRT_OLE; break;
                    default : result = JRT_RE;
                }
                write_runtime_err(strsignal(exitcode));
            }
            ptrace(PTRACE_KILL, pid, NULL, NULL);
            break;
        }
        if (WIFSIGNALED(status)) {
            sig = WTERMSIG(status);
            if (result == JRT_AC) {
                switch (sig) {
                    case SIGCHLD:
                    case SIGALRM: alarm(0);
                    case SIGKILL:
                    case SIGXCPU: result = JRT_TLE; break;
                    case SIGXFSZ: result = JRT_OLE; break;
                    default: result = JRT_RE;
                }
                write_runtime_err(strsignal(sig));
            }
            break;
        }
        // check the system calls
        ptrace(PTRACE_GETREGS, pid, NULL, &reg);

        if (call_counter[reg.REG_SYSCALL] == 0) { //do not limit JVM syscall for using different JVM
            result = JRT_RE;

            char error[BUFFER_SIZE];
            sprintf(error, "A Not allowed system call: solution_id[%d] callid[%ld]\n", solution_id, reg.REG_SYSCALL);
            LOG(log, write_warning, error);
            write_runtime_err(error);
            ptrace(PTRACE_KILL, pid, NULL, NULL);
        } else {
            if (sub == 1 && call_counter[reg.REG_SYSCALL] > 0) --call_counter[reg.REG_SYSCALL];
        }
        sub = 1 - sub;
        ptrace(PTRACE_SYSCALL, pid, NULL, NULL);
    }
    used_time += (ruse.ru_utime.tv_sec * 1000 + ruse.ru_utime.tv_usec / 1000);
    used_time += (ruse.ru_stime.tv_sec * 1000 + ruse.ru_stime.tv_usec / 1000);
}

int
Runner::execute_cmd(const char *fmt, ...)
{
    char cmd[BUFFER_SIZE];
    int ret = 0;
    va_list ap;

    va_start(ap, fmt);
    vsprintf(cmd, fmt, ap);
    ret = system(cmd);
    va_end(ap);
    return ret;
}

void
judge_solution(int &PEflg, double num_of_test)
{
    int comp_res;

    if (result == JRT_AC && used_time > time_lmt * 1000) {
        result = JRT_TLE; 
    }
    if (used_mem > mem_lmt * STD_MB) result = JRT_MLE;

    if (result == JRT_AC) {
        if (spj) {
            comp_res = execute_cmd("%s/data/%d/spj %s %s %s", home_dir.c_str(), problem_id, INPUT_DATA_FILE, OUTPUT_DATA_FILE, USER_OUTPUT_FILE);
            if (comp_res == 0) comp_res = JRT_AC;
            else comp_res = JRT_WA;
        } else {
            comp_res = compare(OUTPUT_DATA_FILE, USER_OUTPUT_FILE);
        }
        result = comp_res;
    }
}

void
Runner::find_next_nonspace(int &c1, int &c2, FILE *&f1, FILE *&f2, JUDGE_RES_TYPE &ret)
{
    while (isspace(c1) || isspace(c2)) {
        if (c1 != c2) {
            if (c2 == EOF) {
                do c1 = fgetc(f1); while (isspace(c1));
                continue;
            } else if (c1 == EOF) {
                do c2 = fgetc(f2); while (isspace(c2));
                continue;
            } else if ((c1 == '\r' && c2 == '\n')) {
                c1 = fgetc(f1);
            } else if ((c2 == '\r' && c1 == '\n')) {
                c2 = fgetc(f2);
            } else {
                ret = JRT_PE;
            }
        }
        if (isspace(c1)) c1 = fgetc(f1);
        if (isspace(c2)) c2 = fgetc(f2);
    }
}

void
Runner::make_diff_out(FILE *&f1, FILE *&f2, int c1, int c2)
{
    FILE *out;
    char buf[45];

    out = fopen(DIFF_OUTPUT_FILE, "a+");
    fprintf(out, "=================\nRight:\n%c", c1);
    if (fgets(buf, 44, f1)) fprintf(out, "%s", buf);
    fprintf(out, "\n-----------------\nYour:\n%c", c2);
    if (fgets(buf, 44, f2)) fprintf(out, "%s", buf);
    fprintf(out, "\n=================\n");
    fclose(out);
}

int
Runner::compare(const char *file1, const char *file2)
{
    JUDGE_RES_TYPE ret = JRT_AC;
    int c1, c2;
    FILE *f1, *f2 ;
    f1 = fopen(file1, "r");
    f2 = fopen(file2, "r");
    
    if (!f1 || !f2) ret = JRT_RE;
    else {
        while (true) {
            c1 = fgetc(f1);
            c2 = fgetc(f2);
            find_next_nonspace(c1, c2, f1, f2, ret);
            while (true) {
                while ((!isspace(c1) && c1) || (!isspace(c2) && c2)) {
                    if (c1 == EOF && c2 == EOF) goto end;
                    if (c1 == EOF || c2 == EOF) break;
                    if (c1 != c2) {
                        ret = JRT_WA;
                        goto end;
                    }
                    c1 = fgetc(f1);
                    c2 = fgetc(f2);
                }
                find_next_nonspace(c1, c2, f1, f2, ret);
                if (c1 == EOF && c2 == EOF) goto end;
                if (c1 == EOF || c2 == EOF) {
                    ret = JRT_WA;
                    goto end;
                }
                if ((c1 == '\n' || !c1) && (c2 == '\n' || !c2)) break;
            }
        }
    }
end:
    if (ret == JRT_WA) make_diff_out(f1, f2, c1, c2);
	if (f1) fclose(f1);
    if (f2) fclose(f2);
    return ret;
}
